# main.py
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, TaskType
import torch

model_name_or_path = "./Qwen2.5-7B-Instruct"  # 本地 Qwen 7B 模型路径
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,
    device_map="auto",
    torch_dtype=torch.float16  # FP16 加速
)

dataset = load_dataset("json", data_files="liucixin.json")
dataset = dataset["train"]

def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512)

tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False  # Causal LM，不做 Masked LM
)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],  # 根据 Qwen 注意力结构选
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)
model = get_peft_model(model, lora_config)

training_args = TrainingArguments(
    output_dir="./qwen-liucixin-lora",
    per_device_train_batch_size=2,   # 显存小可调小
    gradient_accumulation_steps=8,  # 补偿小 batch size
    learning_rate=2e-4,
    num_train_epochs=3,
    logging_steps=50,
    save_steps=500,
    fp16=True,
    save_total_limit=2
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator
)

trainer.train()
trainer.save_model("./qwen-liucixin-lora")

prompt = "在遥远的未来，地球文明面临"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_length=200, do_sample=True, temperature=0.8)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
